import numpy as np
import matplotlib.pyplot as plt
from cloud_detection import *
import cv2
from matplotlib.animation import FuncAnimation

def vis_single_img(time_stamp, key, image):
    # run cloud_detection function
    cloud_cover,cloud_mask,sun_mask, SAMPI = cloud_detection(time_stamp, key, image)
    return cloud_cover, SAMPI
    # sun_center_x, sun_center_y, _ = sun_position(time_stamp)
    
    # f,ax = plt.subplots(1,2)
    # ## original image
    # ax[0].imshow(image[:,:,::-1])
    # #ax[0].set_title(time_stamp)
    # ax[0].set_xticks([])
    # ax[0].set_yticks([])
    # ax[0].set_title('Original image')
    
    # ## cloud detection result
    # ### calculate the morphological gradient for cloud mask to draw the edge to the clouds
    # kernel = np.ones((2,2), np.uint8)
    # bound_cloud = cv2.morphologyEx(cloud_mask, cv2.MORPH_GRADIENT, kernel)

    # ax[1].imshow(image[:,:,::-1], interpolation='none')
    # ax[1].imshow(sun_mask, interpolation='none', alpha=0.15)
    # ax[1].imshow(cloud_mask, interpolation='none', alpha=0.1)
    # ax[1].imshow(bound_cloud, interpolation='none', alpha=0.2)
    # #ax[1].set_title(time_stamp)
    # ax[1].set_xticks([])
    # ax[1].set_yticks([])
    # ax[1].set_title('Cloud detection result') 
    # #ax[1].text(1.1,0.8,'Sun poisition: ({0},{1})'.format(sun_center_x,sun_center_y),transform=ax[1].transAxes)
    # ax[1].text(0.25,0.025,'Cloud fraction={0:.2f}'.format(cloud_cover),color='white',transform=ax[1].transAxes)

    # f.tight_layout()
    # plt.show()

def vis_img_ts(time_stamps,images):
    '''
    Visualization of cloud detection results for image time series
    Generating a video
    '''
    
    # define a images_overlay variable that store original images and cloud detection results
    images_overlay = np.zeros_like(images)
    cloud_cover_arr = np.zeros(len(time_stamps))
    
    # run cloud detection function and overlay results with original images 
    for i in range(len(time_stamps)):
        # run cloud_detection function
        cloud_cover,cloud_mask,sun_mask = cloud_detection(time_stamps[i],images[i])
        cloud_cover_arr[i] = cloud_cover

        # calculate the morphological gradient for cloud mask to draw the edge to the clouds
        kernel = np.ones((2,2), np.uint8)
        bound_cloud = cv2.morphologyEx(cloud_mask, cv2.MORPH_GRADIENT, kernel)

        alpha1=0.15
        images_overlay[i] = cv2.addWeighted(images[i], 1-alpha1, sun_mask[:,:,::-1], alpha1, 0)

        alpha2=0.1
        images_overlay[i] = cv2.addWeighted(images_overlay[i], 1 - alpha2, cloud_mask, alpha2, 0)
        
        alpha3=0.2
        images_overlay[i] = cv2.addWeighted(images_overlay[i], 1 - alpha3, bound_cloud, alpha3, 0)
    
    f, ax = plt.subplots(1,2)
    img_plot1 = ax[0].imshow(images[0,:,:,::-1])
    img_plot2 = ax[1].imshow(images_overlay[0,:,:,::-1])
    cloud_fraction_text = ax[1].text(0.25,0.025, '', color='white', transform=ax[1].transAxes)
    
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    ax[0].set_title('Original image',fontsize=8)
    ax[1].set_xticks([])
    ax[1].set_yticks([])
    ax[1].set_title('Cloud detection result',fontsize=8)

    f.tight_layout()

    def update(frame_idx):
        # Update the image data for each frame
        img_plot1.set_data(images[frame_idx,:,:,::-1])
        img_plot2.set_data(images_overlay[frame_idx,:,:,::-1])

        # Update text elements
        cloud_fraction_text.set_text('Cloud fraction={0:.2f}'.format(cloud_cover_arr[frame_idx]))

        return [img_plot1, img_plot2, cloud_fraction_text]

    # Create the animation
    ani = FuncAnimation(f, update, frames=len(time_stamps), interval=50, blit=True)
    
    plt.close(f)

    return ani

def classify_weather(cloudiness, smapi):
    if cloudiness <= 0.05:
        return 'Sunny'
    elif (0.05 < cloudiness <= 0.8) or (cloudiness > 0.8 and smapi >= 195):
        return 'Cloudy'
    else:
        return 'Overcast'
    
def classify_weather_2(cloudiness, smapi):
    if cloudiness <= 0.05:
        return 'Sunny'
    else:
        return 'Cloudy'